#!/usr/bin/env python


import sys
import PyQt4.Qt as Qt
import PyQt4.Qwt5 as Qwt
import numpy as np


class CandlestickCurve(Qwt.QwtPlotCurve):

    def __init__(self,
                 x, open_ = [], close = [], high = None, low = None,
                 candlePen = Qt.QPen(Qt.Qt.black),
                 upColor = Qt.Qt.white,
                 downColor = Qt.Qt.black
                 ):
        """Candlesticks using open, close, high and low price information.
        
        high and low are optional; if None, only the candle bodies will be
        drawn.

        open_ and close must be sequences with a shape (N,), and
        high and low must be either None or also sequences with a shape (N,).
        
        candlePen is the pen used to draw the candle shadows and body outlines
        
        upColor is the color used to paint the candle bodies when close > open_

        downColor is the color used to paint the candle bodies when 
        open_ > close
        
        """

        Qwt.QwtPlotCurve.__init__(self)
        self.SetData(x, open_, close, high, low)
        #self.setPen(curvePen)
        #self.setStyle(curveStyle)
        #self.setSymbol(curveSymbol)
        self.candlePen = candlePen
        self.upColor = upColor
        self.downColor = downColor

    # __init__()

    def SetData(self, x, open_, close, high=None, low=None):
        """Sets open, close, high (optional) and low (optional) sequences

        high and low are optional; if None, only the candle bodies will be
        drawn.

        open and close must be sequences with a shape (N,), and
        high and low must be either None or also sequences with a shape (N,).

        """
        
        self.__open = np.asarray(open_, float)
        self.__x = x
        if len(self.__x.shape) != 1:
            raise RuntimeError, 'len(asarray(open).shape) != 1'
        self.__close = np.asarray(close, float)
        if len(self.__close.shape) != 1:
            raise RuntimeError, 'len(asarray(close).shape) != 1'
        if len(self.__open) != len(self.__close):
            raise RuntimeError, 'len(asarray(open)) != len(asarray(close))' 

        if high is None:
            self.__high = None
            self.__low = None
        else:
            self.__high = np.asarray(high, float)
            self.__low = np.asarray(low, float)
            if len(self.__high) != len(self.__open):
                raise RuntimeError('len(asarray(high)) != len(asarray(open))')
            if len(self.__low) != len(self.__open):
                raise RuntimeError('len(asarray(low)) != len(asarray(open))')
            
        #Qwt.QwtPlotCurve.setData(self, self.__x, self.__y)

    # setData()
        
    def boundingRect(self):
        """Returns the bounding rectangle of the data.
        
        For the y boundaries, uses low/high if they are used, or open/close
        otherwise.
        
        This function overrides QwtPlotCurve.boundingRect()
        """
        xmin = np.min(self.__x)
        xmax = np.max(self.__x)

        if self.__high is not None:
            ymin = np.min(self.__low)
            ymax = np.max(self.__high)
        else:
            ymin = np.min(np.min(self.__open), np.min(self.__close))
            ymax = np.max(np.max(self.__open), np.max(self.__close))

        return Qt.QRectF(xmin, ymin, xmax-xmin, ymax-ymin)
        
    # boundingRect()

    def drawFromTo(self, painter, xMap, yMap, first, last = -1):
        """Draw an interval of the curve, including the error bars

        painter is the QPainter used to draw the curve

        xMap is the Qwt.QwtDiMap used to map x-values to pixels

        yMap is the Qwt.QwtDiMap used to map y-values to pixels
        
        first is the index of the first data point to draw

        last is the index of the last data point to draw. If last < 0, last
        is transformed to index the last data point
        """

        if last < 0:
            last = self.dataSize() - 1


        painter.save()
        try:
            XTr = xMap.transform # more compact notation
            YTr = yMap.transform
            x = self.__x
            open_ = self.__open
            close = self.__close
            high = self.__high
            low = self.__low
            flagHilo = high is not None
            n, i, = len(x), 0
            # Finds candle width (same unit as x). It must be a fraction of
            # the horizontal point spacing. But the latter may be varible, so
            # finds the minimum distance between x values to use as
            # "horizontal spacing"
            horizSpacing = np.min(np.diff(x))            
            # candle horizontal span will be
            # x[i]-halfWidth until x[i]+halfWidth
            halfWidth = horizSpacing*0.375
            upBrush = Qt.QBrush(self.upColor, Qt.Qt.SolidPattern);
            downBrush = Qt.QBrush(self.downColor, Qt.Qt.SolidPattern);
            painter.setPen(self.candlePen)
            lines = []
            rects = []
                        
            while i < n:
                x1 = XTr(x[i]-halfWidth)
                x2 = XTr(x[i]+halfWidth)

                # Body
                hei = abs(open_[i]-close[i])
                if close[i] >= open_[i]:
                    painter.setBrush(upBrush)
                    bodyTop = YTr(close[i])
                    bodyBottom = YTr(open_[i])
                else:
                    painter.setBrush(downBrush)
                    bodyBottom = YTr(close[i])
                    bodyTop = YTr(open_[i])
                painter.drawRect(Qt.QRect(x1, bodyTop, 1+x2-x1, \
                                      bodyBottom-bodyTop))

                # Shadows
                if flagHilo:
                    xi = XTr(x[i])
                    trHi = YTr(high[i])
                    trLo = YTr(low[i])
                    if trHi < bodyTop: # "<", ">" logic is inverted because the
                                       # variables are converted to screen
                                       # coordinates
                        painter.drawLine(xi, trHi, xi, bodyTop)
                    if trLo > bodyBottom:
                        painter.drawLine(xi, trLo, xi, bodyBottom)
                                            
                i += 1

        finally:
            painter.restore()

    # drawFromTo()

# class ErrorBarPlotCurve


##############################################################################
# Demo

def make():
    # create a plot with a white canvas
    demo = Qwt.QwtPlot(Qwt.QwtText("Candlestick Demonstation"))
    demo.setCanvasBackground(Qt.Qt.white)
    demo.plotLayout().setAlignCanvasToScales(True)
    demo.setAxisScale(Qwt.QwtPlot.xBottom, 0, 10)

    grid = Qwt.QwtPlotGrid()
    grid.attach(demo)
    grid.setPen(Qt.QPen(Qt.Qt.black, 0, Qt.Qt.DotLine))
    
    # Some arbitrary, rather stupid, data
    x = np.arange(1, 10, 1, float)
    open_ = np.asarray([10, 12, 14, 16, 18, 17, 15, 19, 22, 25], float)
    close = np.asarray([11, 13, 15, 17, 18, 16, 18, 18, 23, 28], float)
    high = np.asarray([11.5, 13.5, 15.5, 17.5, 18.5, 16, 18, 18, 24, 29], float)
    low = np.asarray([9, 11, 13, 15, 17, 15, 14, 13, 12, 11], float)

    curve = CandlestickCurve(
        x=x,
        open_=open_,
        close=close,
        high=high,
        low=low
        )
    ''',
    candlePen=Qt.QPen(Qt.Qt.black, 1),
    upColor=Qt.Qt.blue,
    downColor=Qt.Qt.red
    )
    '''
    curve.attach(demo)
    demo.resize(400, 300)
    demo.show()
    return demo

# make()


def main(args):
    app = Qt.QApplication(args)
    demo = make()
    zoomer = Qwt.QwtPlotZoomer(Qwt.QwtPlot.xBottom,
                               Qwt.QwtPlot.yLeft,
                               Qwt.QwtPicker.DragSelection,
                               Qwt.QwtPicker.AlwaysOff,
                               demo.canvas())
    zoomer.setRubberBandPen(Qt.QPen(Qt.Qt.green))
    picker = Qwt.QwtPlotPicker(Qwt.QwtPlot.xBottom,
                               Qwt.QwtPlot.yLeft,
                               Qwt.QwtPicker.NoSelection,
                               Qwt.QwtPlotPicker.CrossRubberBand,
                               Qwt.QwtPicker.AlwaysOn,
                               demo.canvas())
    picker.setTrackerPen(Qt.QPen(Qt.Qt.red))
    sys.exit(app.exec_())

# main()


if __name__ == '__main__':
    main(sys.argv)


